# -*- coding:utf-8 -*-
"""
滑坡分析工具函数 - 依赖TensorFlow和Keras的部分
包含模型预测、结果保存等需要TensorFlow/Keras的函数
"""
# 支持 Keras 3 + PyTorch 后端的 F1-score 指标 与 F1-loss
# 用法：python f1_torchbackend.py
import os
# 强制使用 torch 后端（也可在环境变量里设置）
os.environ["KERAS_BACKEND"] = "torch"

import datetime
import torch
import keras
from keras import ops as KO
from keras.src.backend import backend
from collections import Counter
import numpy as np
import pandas as pd
import geopandas as gpd
from gwlsa_settings import net_params
from utils.general_utils import read_trainValTest_fromCSV, flatten_neighbours_parallel, save_results  # 从general_utils导入通用函数
from utils.general_utils import reshape_dnn_X, reshape_lstm_X, get_onehot
from utils.general_utils import timer
from joblib import Parallel, delayed
import psutil
import pyarrow as pa
import pyarrow.parquet as pq


# scikit-learn的predict返回类别，predict_proba返回概率。
# keras中predict返回的是概率，没有返回类别的函数。与scikit-learn名字有区别。
# 为了使得名称上一致，这里创建一个predcit_proba函数，用于返回预测概率（1D）
def predict_proba(model, x_test, bat_size=1000):
    """
    使用Keras模型预测概率
    
    Parameters:
    -----------
    model : Keras model
        Keras模型
    x_test : array-like
        测试数据
    bat_size : int
        批次大小
        
    Returns:
    --------
    y_proba_positive : 1D array-like
        正类预测概率
    """
    # y_test 2D one-hot
    y_proba_2d = model.predict(x_test, batch_size=bat_size)
    y_proba_positive = y_proba_2d[:,1]
    return y_proba_positive


# scikit-learn的predict返回类别，predict_proba返回概率。
# keras中predict返回的是概率，没有返回类别的函数。与scikit-learn名字有区别。
# 为了使得名称上一致，这里创建一个predcit函数,用于返回预测类别 （1D）
def predict_class(model, x_test, bat_size=1000):
    """
    使用Keras模型预测类别
    
    Parameters:
    -----------
    model : Keras model
        Keras模型
    x_test : array-like
        测试数据
    bat_size : int
        批次大小
        
    Returns:
    --------
    y_pred_class : array-like
        预测类别
    """
    # y_test 2D one-hot
    y_proba_2d = model.predict(x_test, batch_size=bat_size)
    y_pred_class = np.argmax(y_proba_2d, axis=1)
    return y_pred_class


# scikit-learn的predict返回类别，predict_proba返回概率。
# keras中predict返回的是概率，没有返回类别的函数。与scikit-learn名字有区别。
# 为了使得名称上一致，这里创建一个predcit函数,用于返回预测类别 （1D）
def predict_class_prob(model, x_test, bat_size=1000, model_type='keras_dnn_lstm'):
    """
    预测类别和概率
    
    Parameters:
    -----------
    model : model object
        模型对象（Keras或sklearn）
    x_test : array-like
        测试数据
    bat_size : int
        批次大小
    model_type : str
        模型类型，可选'keras_dnn_lstm'或'rf_lr'
        
    Returns:
    --------
    返回值：1一个元组，组成为：（y_pred_class, y_pred_prob）
    y_pred_class : array-like
        预测类别
    y_pred_prob : array-like
        预测概率
    """
    if model_type == 'keras_dnn_lstm':
        y_pred_prob = predict_proba(model, x_test, bat_size=bat_size)
        y_pred_class = predict_class(model, x_test, bat_size=bat_size)
    elif model_type == 'rf_lr':
        y_pred_prob = model.predict_proba(x_test)[:, 1]
        y_pred_class = model.predict(x_test)
    return y_pred_class, y_pred_prob


# 出图时ke使用：预测所有X的目标y（而不只是预测x_test的目标y），将预测结果保存为文件
def save_predY_toCSV(x, model, save_csv_filename, GeoID, bat_size=1000, model_type='keras_dnn_lstm'):
    """
    给定模型和x，预测x对应的Y并保存结果到CSV文件。可以在最终出图时使用。
    
    Parameters:
    -----------
    x : array-like
        输入数据的特征x
    model : model object
        模型对象，可以是keras模型或者是LR和RF模型。
    save_csv_filename : str
        保存的CSV文件路径
    GeoID : array-like
        地理ID，是一个一维数组。
    bat_size : int
        批次大小
    model_type : str
        模型类型，可选'keras_dnn_lstm'或'rf_lr'。默认为'keras_dnn_lstm'
        
    Returns:
    --------
    y_pred_prob : array-like
        预测概率
    y_pred_class : array-like
        预测类别
    """
    y_pred_class, y_pred_prob = predict_class_prob(model, x, bat_size, model_type)
    if save_csv_filename and (GeoID is not None):
        save_results(GeoID, y_pred_class, y_pred_prob, save_csv_filename)
    else:
        raise Exception('参数错误，请设置参数：save_csv_filename和GeoID')
    return y_pred_prob, y_pred_class

def get_proc_nn_data(model_type='dnn',
                     load_from_csv=False,
                     is_shuffle=True,
                     whether_scaled=True,
                     scaled_base='x_train',
                     whether_return_GeoId=False,
                     drop_column=None,
                     data_folder=None,
                     channels=1,
                     flatten_neighbours=False):
    """
    从csv或者pkl加载数据，并将数据拆分成训练、验证和测试集三个部分，根据模型类型进行相应的数据重塑。

    Parameters
    ----------
    model_type : str, optional
        模型类型，可选 'dnn' 或 'lstm'，默认为 'dnn'。决定使用何种数据重塑方式。
    load_from_csv : bool, optional
        是否从csv文件加载数据，如果不从csv加载，则从pkl加载。默认为False。
    is_shuffle : bool, optional
        是否打乱数据，默认为True。
    whether_scaled : bool, optional
        是否缩放数据，默认为True。该选项只对从csv加载数据时可用。从pkl加载时，加载的数据是缩放好的。
    scaled_base : str, optional
        缩放数据的基准，默认以x_train为基准缩放。也可以以整个数据集为基准进行缩放。
    whether_return_GeoId : bool, optional
        是否返回GeoID，默认为False。
    drop_column : str, optional
        需要丢弃的X列，默认为None。
    data_folder : str, optional
        csv文件的存放目录，当load_from_csv为True时可以设置该值。默认为None。为None时使用gnnwr_settings.py中的存放目录。
    channels : int, optional
        重塑数据形状时的通道数，默认为1。
    flatten_neighbours : bool, optional
        是否展平neighbours列，注意必须要有‘neighbours’列，并且load_from_csv=False，即从pkl或parquet文件加载数据，否则会报错。默认为False,不展平。

    Returns
    -------
    tuple
        返回整个数据集、训练集、验证集、测试集、通道数（channels）、以及输入特征的数量（nb_features)。
        如果whether_return_geo_id参数为True，返回的数据集中会包含对应的GeoID。
    """

    # 数据加载部分（两个函数完全相同）
    if load_from_csv:
        data, train_data, val_data, test_data, nb_features = read_trainValTest_fromCSV(
            data_folder=data_folder,
            is_shuffle=is_shuffle,
            whether_scaled=whether_scaled,
            scaled_base=scaled_base,
            drop_columname=drop_column
        )
        x_all, y_r, geoId_r = data[0], data[1], data[2]
        x_train, y_train_r, geoId_train_r = train_data[0], train_data[1], train_data[2]
        x_val, y_val_r, geoId_val_r = val_data[0], val_data[1], val_data[2]
        x_test, y_test_r, geoId_test_r = test_data[0], test_data[1], test_data[2]
        print('Counter(y_train_r):', Counter(y_train_r))
    else:
        train_df, val_df, test_df = load_df_fromfile(
            net_params['data_load_dir'],
            net_params['max_distance'],
            net_params['resolution'],
            file_format=net_params['data_load_format']
        )
        all_df = pd.concat([train_df, val_df, test_df], ignore_index=False)

        x_cols = net_params['x_column_names']
        nb_features = len(x_cols)
        y_col = net_params['y_column_name']
        id_col = net_params['id_column']
        if not flatten_neighbours: #不展平neighbours列（可能是没有这一列）
            x_all, y_r, geoId_r = all_df[x_cols].values, all_df[y_col].values, all_df[id_col].values
            x_train, y_train_r, geoId_train_r = train_df[x_cols].values, train_df[y_col].values, train_df[id_col].values
            x_val, y_val_r, geoId_val_r = val_df[x_cols].values, val_df[y_col].values, val_df[id_col].values
            x_test, y_test_r, geoId_test_r = test_df[x_cols].values, test_df[y_col].values, test_df[id_col].values
        else: # 展平neighbours列
            geoId_r, x_all, y_r = [df.values for df in flatten_neighbours_parallel(all_df, x_cols, y_col, id_col,
                                                                                   neighbours_column='neighbours')]
            geoId_train_r, x_train, y_train_r = [df.values for df in
                                                 flatten_neighbours_parallel(train_df, x_cols, y_col, id_col,
                                                                             neighbours_column='neighbours')]
            geoId_val_r, x_val, y_val_r = [df.values for df in
                                           flatten_neighbours_parallel(val_df, x_cols, y_col, id_col,
                                                                       neighbours_column='neighbours')]
            geoId_test_r, x_test, y_test_r = [df.values for df in
                                              flatten_neighbours_parallel(test_df, x_cols, y_col, id_col,
                                                                          neighbours_column='neighbours')]
            nb_features = x_test.shape[1]       # 展平neighbours列之后，features的数量变了

    # 根据模型类型进行数据重塑
    if model_type.lower() == 'lstm':
        # LSTM数据重塑：将数据塑形为 (samples, channels, nb_features)
        x_all = reshape_lstm_X(x_all, time_steps=channels, nb_features=nb_features)
        x_train = reshape_lstm_X(x_train, time_steps=channels, nb_features=nb_features)
        x_val = reshape_lstm_X(x_val, time_steps=channels, nb_features=nb_features)
        x_test = reshape_lstm_X(x_test, time_steps=channels, nb_features=nb_features)
    else:  # 默认使用DNN数据重塑
        # DNN数据重塑：将数据塑形为 (samples, nb_features, channels)
        x_all = reshape_dnn_X(x_all, nb_features=nb_features, channels=channels)
        x_train = reshape_dnn_X(x_train, nb_features=nb_features, channels=channels)
        x_val = reshape_dnn_X(x_val, nb_features=nb_features, channels=channels)
        x_test = reshape_dnn_X(x_test, nb_features=nb_features, channels=channels)

    # 将y标签独热化（两个函数完全相同）
    y_onehot = get_onehot(y_r)
    y_train_onehot = get_onehot(y_train_r)
    y_val_onehot = get_onehot(y_val_r)
    y_test_onehot = get_onehot(y_test_r)

    # 根据是否返回GeoID决定返回值
    if not whether_return_GeoId:
        return (x_all, y_onehot), (x_train, y_train_onehot), (x_val, y_val_onehot), \
               (x_test, y_test_onehot), y_test_r, channels, nb_features
    else:
        return (x_all, y_onehot, geoId_r), (x_train, y_train_onehot, geoId_train_r), \
               (x_val, y_val_onehot, geoId_val_r), (x_test, y_test_onehot, geoId_test_r), \
               y_test_r, channels, nb_features


# --------------------------------------------------
# 1. F1-score 指标（支持 PyTorch 后端）
# --------------------------------------------------
@keras.saving.register_keras_serializable()
class F1Score(keras.metrics.Metric):
    """
    二分类 F1-score，Keras 3 + PyTorch 后端可用。
    用法：
        model.compile(..., metrics=[F1Score(name='f1')])
    """
    def __init__(self, name="f1", dtype=None, threshold=0.5):
        super().__init__(name=name, dtype=dtype)
        self.threshold = threshold

        # 累积全局统计量
        self.true_pos  = self.add_weight(shape=(), initializer="zeros", name="tp")
        self.false_pos = self.add_weight(shape=(), initializer="zeros", name="fp")
        self.false_neg = self.add_weight(shape=(), initializer="zeros", name="fn")

        # -------- 关键：注册为 nn.Parameter --------
        if keras.config.backend() == "torch":
            import torch.nn as nn
            for v in self.variables:
                v._value = nn.Parameter(v.value)

    def update_state(self, y_true, y_pred, sample_weight=None):
        y_true = KO.cast(y_true, self.dtype)
        y_pred = KO.cast(y_pred >= self.threshold, self.dtype)

        tp = KO.sum(y_true * y_pred)
        fp = KO.sum((1 - y_true) * y_pred)
        fn = KO.sum(y_true * (1 - y_pred))

        self.true_pos.assign_add(tp)
        self.false_pos.assign_add(fp)
        self.false_neg.assign_add(fn)

    def result(self):
        eps = keras.backend.epsilon()
        precision = self.true_pos / (self.true_pos + self.false_pos + eps)
        recall = self.true_pos / (self.true_pos + self.false_neg + eps)
        f1 = 2 * precision * recall / (precision + recall + eps)
        f1_safe = KO.where(KO.isnan(f1), 0.0, f1)
        # 根据当前后端选择操作
        current_backend = backend()
        if current_backend == "torch":
            import torch
            # 如果 f1_safe 是 Keras 张量，可能需要先 .value 再 .detach()
            # 具体取决于 Keras 版本的封装方式
            result_tensor = f1_safe.detach().clone()
            return result_tensor
        else:
            # 对于其他后端，使用 stop_gradient
            return KO.stop_gradient(f1_safe)

    def reset_state(self):
        # 兼容 Keras 3 所有后端
        for v in self.variables:
            v.assign(KO.zeros_like(v))

# --------------------------------------------------
# 2. F1-loss（无状态，直接可用）
# --------------------------------------------------
@keras.saving.register_keras_serializable()
def f1_loss(y_true, y_pred, eps=None):
    """
    多分类兼容的 F1-loss（宏平均）。
    y_true : one-hot  [B, C]
    y_pred : prob     [B, C]
    """
    eps = eps or keras.backend.epsilon()
    y_true = KO.cast(y_true, y_pred.dtype)

    # 把样本维度规约掉，得到每个类别的 tp/fp/fn
    tp = KO.sum(y_true * y_pred, axis=0)      # [C]
    fp = KO.sum((1 - y_true) * y_pred, axis=0)
    fn = KO.sum(y_true * (1 - y_pred), axis=0)

    precision = tp / (tp + fp + eps)          # [C]
    recall    = tp / (tp + fn + eps)
    f1_per_class = 2 * precision * recall / (precision + recall + eps)
    # 宏平均
    macro_f1 = KO.mean(KO.where(KO.isnan(f1_per_class), 0.0, f1_per_class))

    return 1 - macro_f1

def nestedDataframes_to_dictList(gdf):
    '''
    将DataFrame嵌套DataFrame的形式转换为DataFrame嵌套字典列表的形式
    Parameters
    ----------
    gdf：DataFrame
        待转换的某列元素嵌套子DataFrame的DataFrame

    Returns：DataFrame
    -------
        子DataFrame转换为字典列表后的DataFrame
    '''
    df = gdf.copy()
    # 遍历每一列
    for col in df.columns:
        # 检查列中的元素是否为DataFrame
        if any(isinstance(item, pd.DataFrame) for item in df[col]):
            # 将DataFrame转换为字典列表
            df[col] = df[col].apply(lambda x: x.to_dict('records') if isinstance(x, pd.DataFrame) else x)
    return df

def _process_a_chunk(data_chunk, nested_columns):
    """
    处理一个数据块的函数。这个块包含所有列，但只是所有行的一部分。
    """
    chunk_copy = data_chunk.copy()
    for col in nested_columns:
        chunk_copy[col] = chunk_copy[col].apply(
            lambda x: pd.DataFrame(x.tolist()) if isinstance(x, np.ndarray) and len(x) > 0 and isinstance(x[0], dict) else x
        )
    return chunk_copy
def restore_nested_dataframes_chunked(df, nested_columns, n_jobs=-1, chunksize=None):
    """
    将DataFrame嵌套字典列表的形式转换为DataFrame嵌套DataFrame的形式。
    通过数据分块进行并行处理。
    chunksize: 每个块的行数。如果为None，则根据CPU核心数和数据大小自动估算。
    """
    if chunksize is None:
        logic_workers = psutil.cpu_count(logical=True)  # 逻辑处理器总数
        # 一个简单的启发式估算：将数据分成大约 (CPU核心数 * 4) 个块
        n_chunks = max(1, len(df) // ( int(logic_workers * 0.8) * 4))
        chunksize = max(1, len(df) // n_chunks)

    # 将DataFrame分成多个块
    chunks = [df[i:i + chunksize] for i in range(0, len(df), chunksize)]
    # print(f"将数据分成了 {len(chunks)} 个块进行并行处理。")

    # 并行处理每个块
    processed_chunks = Parallel(n_jobs=n_jobs)(
        delayed(_process_a_chunk)(chunk, nested_columns) for chunk in chunks
    )

    # 将所有处理后的块拼接起来
    df_restored = pd.concat(processed_chunks, ignore_index=True)
    return df_restored

def read_gpkg_addColumns(gpkg_path,
                         id_column='id',
                         layer_name='data',
                         csv_encoding='gb2312',
                         added_columns=['distances', 'neighbours']):
    '''
    从gpkg文件读取数据，并添加指定的列。
    Parameters
    ----------
    gpkg_path:str
        gpkg文件路径
    id_column: str
        gpkg文件中id列的列名，默认为'id'
    layer_name: str
        gpkg文件中要读取的图层名，默认为'data'
    csv_encoding: str
        gpkg文件的编码格式，默认为'gb2312'
    added_columns: list
        要添加的列名列表，默认为['distances', 'neighbours']

    Returns
    -------
    返回一个元组，组成形式：max_id_value, df
    max_id_value: int
        id列的最大值
    df: DataFrame
        添加指定列后的DataFrame

    '''
    if not gpkg_path.endswith('.gpkg'):
        raise ValueError('读取的原始文件必须是gpkg!')
    df = gpd.read_file(gpkg_path, layer=layer_name, encoding=csv_encoding)
    max_id_value = df[id_column].max()

    # 添加'distances', 'neighbours'两列
    for added_column in added_columns:
        df[added_column]=None
    return max_id_value, df

@timer
def load_df_fromfile(file_dir, buffer_distance=91, resolution=90, only_read_test=False, file_format='pkl'):
    """
        从pkl文件加载训练、验证、测试数据集。
    Parameters
    ----------
    file_dir: str
        3个文件所在的目录
    buffer_distance: int
        缓冲区距离，生成文件时使用的距离，用于计算邻居数。默认为91m
    resolution: int
        数据的空间分辨率。默认为90m
    only_read_test: bool
        是否只读取测试数据。默认为False，读取所有数据。当只读取测试数据评估测试精度时可设置为True。
    file_format: str
        读取的训练、验证和测试集这3个文件的后缀名，可以为parquet或者pkl格式。默认为pkl。
        如果是DataFrame中包含嵌入的其他DataFrame，例如包含邻居数据列，那么file_format请使用pkl格式，不要使用parquet格式。

    Returns
    -------
    list
        返回train_df, val_df, test_df或者None, None, test_df（只读取测试数据时）
    """
    if not os.path.exists(file_dir):
        raise FileExistsError(f'目录{file_dir}不存在！！！')
    train_path = os.path.join(file_dir, f'D{buffer_distance}res{resolution}_train.{file_format}')
    val_path = os.path.join(file_dir, f'D{buffer_distance}res{resolution}_val.{file_format}')
    test_path = os.path.join(file_dir, f'D{buffer_distance}res{resolution}_test.{file_format}')

    # 确定是否需要读取训练集和验证集
    read_train_val = not only_read_test
    # 根据文件格式读取数据
    if file_format == 'pkl':
        if read_train_val:
            train_df = pd.read_pickle(train_path)
            val_df = pd.read_pickle(val_path)
        test_df = pd.read_pickle(test_path)
    elif file_format == 'parquet':
        # if read_train_val:
        #     train_df = pd.read_parquet(train_path)
        #     val_df = pd.read_parquet(val_path)
        # test_df = pd.read_parquet(test_path)
        if read_train_val:
            train_table = pq.read_table(train_path)
            train_df = train_table.to_pandas()

            val_table = pq.read_table(val_path)
            val_df = val_table.to_pandas()

        test_table = pq.read_table(test_path)
        test_df = test_table.to_pandas()
        train_len = len(train_df)
        val_len = len(val_df)
        con_df = pd.concat([train_df, val_df, test_df], axis=0, ignore_index=True)
        nested_columns = ['neighbours', 'distances']    # 这两列的元素是嵌套的子DataFrame
        restored_con_df = restore_nested_dataframes_chunked(con_df, nested_columns=nested_columns)

        # 注意需要重置index，因为pytorch的dataset中的get_item方法需要index从0开始
        train_df = restored_con_df.iloc[:train_len].reset_index(drop=True)                      # 前train_len行
        val_df = restored_con_df.iloc[train_len: train_len + val_len].reset_index(drop=True)    # 紧接着的val_len行
        test_df = restored_con_df.iloc[train_len + val_len:].reset_index(drop=True)             # 剩余的所有行

    # 返回结果
    if read_train_val:
        return train_df, val_df, test_df
    else:
        return None, None, test_df

@timer
def save_df_tofile(saved_dir, layer_name, train_df, val_df, test_df, buffer_distance, resolution, file_format, parquet_schema=None):
    '''
        将DataFrame保存到文件。如果DataFrame中包含嵌入的其他DataFrame，那么file_format请使用pkl格式，不要使用parquet格式。
    Parameters
    ----------
    saved_dir: str
        保存目录
    layer_name: str
        图层的名称
    train_df: DataFrame
        训练数据集
    val_df: DataFrame
        验证数据集
    test_df: DataFrame
        测试数据集
    buffer_distance: int
        缓冲区距离，生成文件时使用的距离，用于计算邻居数。
    resolution: int
        数据的空间分辨率。
    file_format: string
        保存文件的格式，可以为parquet或者pkl格式。如果是DataFrame中包含嵌入的其他DataFrame，那么file_format请使用pkl格式，不要使用parquet格式。
    parquet_schema: schema
        当保存为parquet格式时使用的schema。当DataFrame中包含嵌入的其他DataFrame时，必须设置，否则保存会报错。默认值为None。

    Returns
    -------

    '''
    out_folder_name = datetime.datetime.today().strftime("%Y%m%d_%H%M%S")
    out_folder_name = layer_name + '_' + out_folder_name + f'_D{buffer_distance}res{resolution}'
    out_folder_path = os.path.join(saved_dir, out_folder_name)
    if not os.path.exists(out_folder_path):
        os.makedirs(out_folder_path)
    train_path = os.path.join(saved_dir, out_folder_name, f'D{buffer_distance}res{resolution}_train.{file_format}')
    val_path = os.path.join(saved_dir, out_folder_name, f'D{buffer_distance}res{resolution}_val.{file_format}')
    test_path = os.path.join(saved_dir, out_folder_name, f'D{buffer_distance}res{resolution}_test.{file_format}')
    if file_format == 'pkl':
        train_df.to_pickle(train_path)
        val_df.to_pickle(val_path)
        test_df.to_pickle(test_path)
    elif file_format == 'parquet':
        train_table = pa.Table.from_pandas(train_df, schema=parquet_schema)
        val_table = pa.Table.from_pandas(val_df, schema=parquet_schema)
        test_table = pa.Table.from_pandas(test_df, schema=parquet_schema)
        pq.write_table(train_table, train_path)
        pq.write_table(val_table, val_path)
        pq.write_table(test_table, test_path)
        # train_df.to_parquet(train_path, engine='pyarrow')
        # val_df.to_parquet(val_path, engine='pyarrow')
        # test_df.to_parquet(test_path, engine='pyarrow')

# --------------------------------------------------
# 3. Keras 3的可用性快速冒烟测试
# --------------------------------------------------
if __name__ == "__main__":
    # 造点假数据
    x = np.random.randn(2000, 10).astype("float32")
    y = (np.random.rand(2000) > 0.7).astype("float32")

    # 搭个简单模型
    model = keras.Sequential([
        keras.layers.Dense(64, activation="relu"),
        keras.layers.Dense(1,  activation="sigmoid")
    ])

    model.compile(optimizer="adam",
                  loss=f1_loss,
                  metrics=[F1Score()])

    # 训练
    model.fit(x, y,
              validation_split=0.2,
              epochs=5,
              batch_size=256)

    # 若 GPU 可用，可再验证设备一致性
    if torch.cuda.is_available():
        model.cuda()
        x_gpu = torch.tensor(x).cuda()
        y_gpu = torch.tensor(y).cuda()
        print("CUDA F1:", model.evaluate(x_gpu, y_gpu, verbose=0))
